#define LINALTER_SOURCE
#include "lib.h"

int traceme(int y)
{
	int ret = 0;
	spin_lock(&tracedlock);
	if(y){
		if(tracedpid==0){
			tracedpid = current->pid;
			PRINTK("<3>linalter: Traced pid: %ld\n", (long)tracedpid);
		}else if(find_task_by_vpid(tracedpid)==0){
			tracedpid = current->pid;
			PRINTK("<3>linalter: No longer trace killed process. Traced pid: %ld\n", (long)tracedpid);
		}else{
			PRINTK("<3>linalter: Want to trace current process. But already traced another process.\n");
			ret = -1;
		}
	}else{
	 	if(tracedpid==current->pid){
			PRINTK("<3>linalter: Untraced pid: %ld.\n", (long)tracedpid);
			tracedpid = 0;
		}else{
			PRINTK("<3>linalter: current->pid: %ld tracedpid:%ld\nWant to untrace current process. But current process not traced.\n", (long)current->pid, (long)tracedpid );
			ret = -1;
		}
	}
	spin_unlock(&tracedlock);
	return ret;
}

int amitraced(void)
{
	int ret = 0;
	struct task_struct *t;
	
	spin_lock(&tracedlock);
	if(tracedpid == current->pid)
		ret=1;
	else{
		t = current->real_parent;
		do{
			if(tracedpid == t->pid)
				ret=1;
			t = t->real_parent;
		}while(t != t->real_parent);
	}
	spin_unlock(&tracedlock);
	return ret;
}

static int istraced(struct task_struct *r)
{
	int ret = 0;
	struct task_struct *t;
	spin_lock(&tracedlock);
	if(r){
		if(tracedpid == r->pid)
			ret=1;
		else{
			t = r->real_parent;
			do{
				if(tracedpid == t->pid)
					ret=1;
				t = t->real_parent;
			}while(t != t->real_parent);
		}
	}
	spin_unlock(&tracedlock);
	return ret;
}

struct task_struct* tracedtask(void)
{
	struct task_struct *r=0;
	spin_lock(&tracedlock);
	if(tracedpid)
		r = find_task_by_vpid(tracedpid);
	spin_unlock(&tracedlock);
	if(r){
		if(tracedpid == r->pid)
			return r;
		else
			PRINTK("<3>linalter: BUG! tracedpid %ld != r->pid %ld\n", (long)tracedpid, (long)r->pid);
	}else
		PRINTK("<3>linalter: traced process not exists.\n");
	return 0;
}

void showpid(struct task_struct *t, char *page, int *len)
{
	struct task_struct *child;
	list_for_each_entry(child, &t->children, sibling) {
		*len += sprintf(page + *len, "%ld\n", (long)child->pid);
		showpid(child, page, len);
	}
}

static int check1(const char *from, unsigned long len, const char *to)
{
	if(len > 128*1024*1024)
		return -LEN_TOO_LONG;
	if( ((unsigned long)from & (PAGE_SIZE-1))!=0 || (len & (PAGE_SIZE-1))!=0 )
		return -ADDR_NOT_ALIGNED;
	if(!access_ok(VERIFY_WRITE, to, len + (len+PAGE_SIZE-1)/PAGE_SIZE) )
		return -CANNOT_WRITE_TO;
	if((unsigned long)from + len >= PAGE_OFFSET)
		return -READ_KERNEL;
	return 0;
}

static int check2(struct task_struct *t)
{
	if(t == 0)
		return -NO_TRACED_TASK;
	if(t->pid == current->pid)
		return -TRY_COPY_SELF;
	if(t->state != TASK_STOPPED)
		return -TASK_NOT_STOPPED;
	if(t->uid != current->uid && current->uid != 0)
		return -NO_PRIVILEGE;
	if(!istraced(t))
		return -TASK_NOT_TRACED;
	return 0;	
}

static unsigned long to_pa(struct mm_struct *mm, unsigned long va)
{
	unsigned long pa;
	pgd_t *pgd;
	pgd = pgd_offset(mm, va);
	if( pgd_present(*pgd) ){
		pud_t *pud = pud_offset(pgd, va);
		if( pud_present(*pud) ){
			pmd_t *pmd = pmd_offset(pud, va);
			if( pmd_present(*pmd) ){
				pte_t *pte = pte_offset_kernel(pmd, va);
				if( pte_present(*pte) ){
					pa = __pa(page_address(pte_page(*pte)));
					pa |= (va & (PAGE_SIZE-1) );
					return pa;
				}
			}
		}
	}
	return 0;
}

static void do_copy(struct mm_struct *mm, char *from, unsigned long len, char *to)
{
	char *success = to + len; /* Use every byte to represent whether a page is successfully copied.*/
	unsigned long va, pa;
	
	spin_lock(&mm->page_table_lock);
	for(va = (unsigned long)from; va < (unsigned long)from+len; va += PAGE_SIZE)
	{
		pa = to_pa(mm, va);
		*(success++) = (pa!=0?1:0);
		if(pa){
			copy_to_user(to, __va(pa), PAGE_SIZE);
		}
		to += PAGE_SIZE;
	}
	spin_unlock(&mm->page_table_lock);
}

int copy(pid_t pid, char *from, unsigned long len, char *to)
{
	struct task_struct *t;
	struct mm_struct *mm;
	struct vm_area_struct *vma;
	int found = 0;
	int ret;
	
	ret = check1(from, len, to);
	if(ret)
		return ret;
	t = find_task_by_vpid(pid);
	ret = check2(t);
	if(ret)
		return ret;
	
	mm = get_task_mm(t);
	if(mm == 0)
		return -NO_MM;
	down_read(&mm->mmap_sem);
	for (vma = mm->mmap; vma; vma = vma->vm_next) {
		if( (vma->vm_flags & VM_WRITE) && 
		  vma->vm_start <= (unsigned long)from &&
		  (unsigned long)from + len <= vma->vm_end){
			found=1;
			do_copy(mm, from, len, to);
			break;
		}
	}
	up_read(&mm->mmap_sem);
	mmput(mm);
	if(!found)
		return -NO_VMA;
	else
		return 0;
}

static int do_write_mem(struct mm_struct *mm, char *to, char *from, unsigned long len)
{
	int count;
	unsigned long pa;
	int ret = 0;

	spin_lock(&mm->page_table_lock);
	while(len>0){
		pa = to_pa(mm, (unsigned long)to);
		if(pa == 0){
			ret = -NO_PA;
			break;
		}
		count = min(len, PAGE_SIZE - (pa & (PAGE_SIZE-1)) );
		copy_from_user(__va(pa), from, count);
		len -= count;
		from += count;
		to += count;
	}
	spin_unlock(&mm->page_table_lock);
	return ret;
}

int write_mem(pid_t pid, char *to, char *from, unsigned long len)
{
	struct task_struct *t;
	struct mm_struct *mm;
	struct vm_area_struct *vma;
	int found = 0;

	if(len > 50*1024*1024)
		return -LEN_TOO_LONG;
	if(!access_ok(VERIFY_READ, from, len)) 
		return -CANNOT_READ_FROM;
	if((unsigned long)from + len >= PAGE_OFFSET)
		return -READ_KERNEL;
	t = find_task_by_vpid(pid);
	if(t == 0)
		return -NO_TRACED_TASK;
	if(t->pid == current->pid)
		return -TRY_COPY_SELF;
	if(t->state != TASK_STOPPED)
		return -TASK_NOT_STOPPED;
	if(t->uid != current->uid && current->uid != 0)
		return -NO_PRIVILEGE;
	if(!istraced(t))
		return -TASK_NOT_TRACED;
	
	mm = get_task_mm(t);
	if(mm == 0)
		return -NO_MM;
	down_read(&mm->mmap_sem);
	for (vma = mm->mmap; vma; vma = vma->vm_next) {
		if( (vma->vm_flags & VM_WRITE) && 
		  vma->vm_start <= (unsigned long)to &&
		  (unsigned long)to + len <= vma->vm_end){
			found=1;
			do_write_mem(mm, to, from, len);
			break;
		}
	}
	up_read(&mm->mmap_sem);
	mmput(mm);
	if(!found)
		return -NO_VMA;
	else
		return 0;
}
